from flask import Flask, request, jsonify, Response
import requests
import json

# 创建Flask应用实例
app = Flask(__name__)

# 默认API配置
DEFAULT_URL = "https://api.apiyi.com/v1/chat/completions"
DEFAULT_MODEL = "gpt-5-mini"
DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."

# API密钥（注意：在生产环境中应使用环境变量存储）
API_KEY = "sk-pEvO0Ks4uHYbvFrE2aA387Eb93534cDb815878D2F3D5B1D2"

@app.route('/chat/completions', methods=['POST'])
def chat_completions():
    """
    处理聊天完成请求的接口
    
    请求体应包含:
    - messages: 消息列表
    - model: (可选) 模型名称
    - stream: (可选) 是否使用流式响应
    
    返回:
    - JSON格式的API响应
    """
    try:
        # 获取请求数据
        request_data = request.get_json()
        
        # 如果没有提供消息，则返回错误
        if not request_data or 'messages' not in request_data:
            return jsonify({
                'error': 'Missing messages in request body'
            }), 400
        
        # 构建请求参数
        url = DEFAULT_URL
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        }
        
        # 准备发送给API的数据
        data = {
            "model": request_data.get("model", DEFAULT_MODEL),
            "stream": request_data.get("stream", False),
            "messages": request_data["messages"]
        }
        
        # 如果请求中包含system prompt，则使用它替换默认的
        if len(data["messages"]) > 0 and data["messages"][0]["role"] != "system":
            data["messages"].insert(0, {
                "role": "system",
                "content": DEFAULT_SYSTEM_PROMPT
            })
        
        # 发送请求到AI API
        response = requests.post(url, headers=headers, json=data)
        
        # 处理响应
        if response.status_code == 200:
            # 如果是流式响应，直接返回流
            if data.get("stream", False):
                return Response(
                    response.content,
                    content_type='text/event-stream',
                    status=200
                )
            else:
                # 非流式响应，返回JSON
                return jsonify(response.json())
        else:
            # API返回错误，转发错误信息
            return jsonify({
                'error': f"API Error: {response.status_code}",
                'message': response.text
            }), response.status_code
            
    except Exception as e:
        # 处理其他异常
        return jsonify({
            'error': 'Internal Server Error',
            'message': str(e)
        }), 500

@app.route('/health', methods=['GET'])
def health_check():
    """
    健康检查接口
    
    Returns:
        JSON格式的健康状态信息
    """
    return jsonify({
        'status': 'healthy',
        'message': 'AI service is running'
    })

@app.route('/', methods=['GET'])
def home():
    """
    主页接口，提供API使用说明
    
    Returns:
        JSON格式的API文档信息
    """
    return jsonify({
        'message': 'AI Service API',
        'endpoints': {
            '/chat/completions': 'POST - Chat completions endpoint',
            '/health': 'GET - Health check endpoint',
            '/': 'GET - This help message'
        },
        'usage': {
            'method': 'POST',
            'url': '/chat/completions',
            'headers': {
                'Content-Type': 'application/json'
            },
            'body': {
                'messages': [
                    {'role': 'user', 'content': 'Hello!'}
                ]
            }
        }
    })

if __name__ == '__main__':
    # 启动Flask应用
    # debug=False 用于生产环境，True用于开发环境
    app.run(host='0.0.0.0', port=5000, debug=True)
